#pragma once

#include <linux/rbtree.h>

/*
 * Please note - only struct rb_augment_callbacks and the prototypes for
 * rb_insert_augmented() and rb_erase_augmented() are intended to be public.
 * The rest are implementation details you are not expected to depend on.
 *
 * See Documentation/core-api/rbtree.rst for documentation and samples.
 */

struct rb_augment_callbacks
{
    void (*propagate)(struct rb_node *node, struct rb_node *stop);
    void (*copy)(struct rb_node *old, struct rb_node *new);
    void (*rotate)(struct rb_node *old, struct rb_node *new);
};

extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
                                  void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
extern void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
                             void (*augment_rotate)(struct rb_node *old, struct rb_node *new));

struct rb_node *rb_add_augmented_cached(struct rb_node *node, struct rb_root_cached *tree,
                                        bool (*less)(struct rb_node *, const struct rb_node *),
                                        const struct rb_augment_callbacks *augment);

#define RB_RED 0
#define RB_BLACK 1

#define __rb_parent(pc) ((struct rb_node *)(pc & ~3))

#define __rb_color(pc) ((pc) & 1)
#define __rb_is_black(pc) __rb_color(pc)
#define __rb_is_red(pc) (!__rb_color(pc))
#define rb_color(rb) __rb_color((rb)->__rb_parent_color)
#define rb_is_red(rb) __rb_is_red((rb)->__rb_parent_color)
#define rb_is_black(rb) __rb_is_black((rb)->__rb_parent_color)

/*
 * Template for declaring augmented rbtree callbacks (generic case)
 *
 * RBSTATIC:    'static' or empty
 * RBNAME:      name of the rb_augment_callbacks structure
 * RBSTRUCT:    struct type of the tree nodes
 * RBFIELD:     name of struct rb_node field within RBSTRUCT
 * RBAUGMENTED: name of field within RBSTRUCT holding data for subtree
 * RBCOMPUTE:   name of function that recomputes the RBAUGMENTED data
 */

#define RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME,                          \
                             RBSTRUCT, RBFIELD, RBAUGMENTED, RBCOMPUTE) \
    static inline void                                                  \
    RBNAME##_propagate(struct rb_node *rb, struct rb_node *stop)        \
    {                                                                   \
        while (rb != stop)                                              \
        {                                                               \
            RBSTRUCT *node = rb_entry(rb, RBSTRUCT, RBFIELD);           \
            if (RBCOMPUTE(node, true))                                  \
                break;                                                  \
            rb = rb_parent(&node->RBFIELD);                             \
        }                                                               \
    }                                                                   \
    static inline void                                                  \
    RBNAME##_copy(struct rb_node *rb_old, struct rb_node *rb_new)       \
    {                                                                   \
        RBSTRUCT *old = rb_entry(rb_old, RBSTRUCT, RBFIELD);            \
        RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD);            \
        new->RBAUGMENTED = old->RBAUGMENTED;                            \
    }                                                                   \
    static void                                                         \
    RBNAME##_rotate(struct rb_node *rb_old, struct rb_node *rb_new)     \
    {                                                                   \
        RBSTRUCT *old = rb_entry(rb_old, RBSTRUCT, RBFIELD);            \
        RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD);            \
        new->RBAUGMENTED = old->RBAUGMENTED;                            \
        RBCOMPUTE(old, false);                                          \
    }                                                                   \
    RBSTATIC const struct rb_augment_callbacks RBNAME = {               \
        .propagate = RBNAME##_propagate,                                \
        .copy = RBNAME##_copy,                                          \
        .rotate = RBNAME##_rotate};

struct rb_node *__rb_erase_augmented(struct rb_node *node, struct rb_root *root,
                                     const struct rb_augment_callbacks *augment);
void __rb_change_child(struct rb_node *old, struct rb_node *new,
                       struct rb_node *parent, struct rb_root *root);

static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
{
    rb->__rb_parent_color = rb_color(rb) + (uintptr_t)p;
}

static inline void rb_set_parent_color(struct rb_node *rb,
                                       struct rb_node *p, int color)
{
    rb->__rb_parent_color = (uintptr_t)p + color;
}

static inline void rb_erase_augmented(struct rb_node *node, struct rb_root *root,
                                      const struct rb_augment_callbacks *augment)
{
    struct rb_node *rebalance = __rb_erase_augmented(node, root, augment);
    if (rebalance)
        __rb_erase_color(rebalance, root, augment->rotate);
}

static inline void rb_erase_augmented_cached(struct rb_node *node,
                                             struct rb_root_cached *root,
                                             const struct rb_augment_callbacks *augment)
{
    if (root->rb_leftmost == node)
        root->rb_leftmost = rb_next(node);
    rb_erase_augmented(node, &root->rb_root, augment);
}

/*
 * Template for declaring augmented rbtree callbacks,
 * computing RBAUGMENTED scalar as max(RBCOMPUTE(node)) for all subtree nodes.
 *
 * RBSTATIC:    'static' or empty
 * RBNAME:      name of the rb_augment_callbacks structure
 * RBSTRUCT:    struct type of the tree nodes
 * RBFIELD:     name of struct rb_node field within RBSTRUCT
 * RBTYPE:      type of the RBAUGMENTED field
 * RBAUGMENTED: name of RBTYPE field within RBSTRUCT holding data for subtree
 * RBCOMPUTE:   name of function that returns the per-node RBTYPE scalar
 */

#define RB_DECLARE_CALLBACKS_MAX(RBSTATIC, RBNAME, RBSTRUCT, RBFIELD,    \
                                 RBTYPE, RBAUGMENTED, RBCOMPUTE)         \
    static inline bool RBNAME##_compute_max(RBSTRUCT *node, bool exit)   \
    {                                                                    \
        RBSTRUCT *child;                                                 \
        RBTYPE max = RBCOMPUTE(node);                                    \
        if (node->RBFIELD.rb_left)                                       \
        {                                                                \
            child = rb_entry(node->RBFIELD.rb_left, RBSTRUCT, RBFIELD);  \
            if (child->RBAUGMENTED > max)                                \
                max = child->RBAUGMENTED;                                \
        }                                                                \
        if (node->RBFIELD.rb_right)                                      \
        {                                                                \
            child = rb_entry(node->RBFIELD.rb_right, RBSTRUCT, RBFIELD); \
            if (child->RBAUGMENTED > max)                                \
                max = child->RBAUGMENTED;                                \
        }                                                                \
        if (exit && node->RBAUGMENTED == max)                            \
            return true;                                                 \
        node->RBAUGMENTED = max;                                         \
        return false;                                                    \
    }                                                                    \
    RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME,                               \
                         RBSTRUCT, RBFIELD, RBAUGMENTED, RBNAME##_compute_max)

static inline void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
                    const struct rb_augment_callbacks *augment)
{
    __rb_insert_augmented(node, root, augment->rotate);
}
